package com.sdy.dcsb.web.filter;

import org.springframework.stereotype.Component;
import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.IOException;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * header中增加req_type
 */
@Component
public class AjaxFilter implements Filter {
    
    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
        HttpServletRequest request = (HttpServletRequest) servletRequest;
        ModifyHttpServletRequestWrapper httpServletRequestWrapper = new ModifyHttpServletRequestWrapper(request);
        httpServletRequestWrapper.putHeader("req_type", "async");
        filterChain.doFilter(httpServletRequestWrapper, servletResponse);
    }

    public class ModifyHttpServletRequestWrapper extends HttpServletRequestWrapper {
        private Map<String, String> customHeaders;

        ModifyHttpServletRequestWrapper(HttpServletRequest request) {
            super(request);
            this.customHeaders = new HashMap<>();
        }

        void putHeader(String name, String value) {
            this.customHeaders.put(name, value);
        }

        @Override
        public String getHeader(String name) {
            String value = this.customHeaders.get(name);
            if (value != null) {
                return value;
            }
            return ((HttpServletRequest) getRequest()).getHeader(name);
        }

        @Override
        public Enumeration<String> getHeaderNames() {
            Set<String> set = new HashSet<>(customHeaders.keySet());
            Enumeration<String> enumeration = ((HttpServletRequest) getRequest()).getHeaderNames();
            while (enumeration.hasMoreElements()) {
                String name = enumeration.nextElement();
                set.add(name);
            }
            return Collections.enumeration(set);
        }
    }
}
